from numpy import *
from numpy.random import uniform
from scipy import *

from quaternion import *

class Sampling:
    def __init__(self, ini, v_ini):
        assert len(ini.shape) == 1
        #assert ini.shape == v_ini.shape

        self.samples = r_[[ini]]
        self.Ns = 1
        self.Ndim = ini.shape[0]
        self.weights = ones(1)
        self.weights_cs = cumsum(self.weights)
        self.weights_cs = self.weights_cs / self.weights_cs[-1]
        
        self.variance = v_ini
        self.beta = 0.01

    def generate_posterior(self, new_Ns):
        new_samples = zeros((new_Ns, self.Ndim))
        for n in range(new_Ns):
            r = uniform()
            s = nonzero(self.weights_cs >= r)[0][0]
            new_s = self.samples[s] + uniform(-self.variance, self.variance)
            new_s[0] = (new_s[0])%1.0
            new_samples[n] = new_s

        self.samples = new_samples
        self.mlogL = ones(new_Ns)
        self.weights = ones(new_Ns)
        self.Ns = new_Ns

    def generate_posterior_quaternion(self, new_Ns):
        new_samples = zeros((new_Ns, self.Ndim))
        for n in range(new_Ns):
            r = uniform()
            s = nonzero(self.weights_cs >= r)[0][0]

            new_s = random_quaternion(self.variance[0])
            new_s = new_s * Quat(self.samples[s])
            
            new_s.normalize()
            new_samples[n] = new_s.q

        self.samples = new_samples
        self.mlogL = ones(new_Ns)
        self.weights = ones(new_Ns)
        self.Ns = new_Ns

    def modulate_samples(self, func):
        for n in range(self.Ns):
            mlogL = func(self.samples[n])
            self.mlogL[n] = mlogL
            self.weights[n] = exp(self.beta *(-mlogL+2000))
        self.weights_cs = cumsum(self.weights)

        if self.weights_cs[-1] == 0:
            raise 'All samples are bad, all weights == 0.'
        self.weights_cs = self.weights_cs / self.weights_cs[-1]

